# 更改预测过程中的is_training参数为True
# Date: 2021/11/5
# Usage: 添加计算测试准确度的程序（dice-loss）
# Modify Date: 2021/11/18
# Usage: 在输出预测为512的尺寸上与标签比对
import cv2
import numpy as np
import tensorflow as tf
import logging
log = logging.getLogger(__name__)

import matplotlib.pyplot as plt

import model_temp as model
# import model as model
import time

import os
import random

class Detector(object):
    def __init__(self,model_dir):
        os.environ['CUDA_VISIBLE_DEVICES'] = '1'
        config = tf.ConfigProto(allow_soft_placement=True)
        config.gpu_options.per_process_gpu_memory_fraction = 1.0
        config.gpu_options.allow_growth = True
        self.input_images = tf.placeholder(tf.float32, shape=[None, None, None, 3], name='input_images')
        self.session = tf.Session(config=config)
        self.global_step = tf.get_variable('global_step', [], initializer=tf.constant_initializer(0), trainable=False)
        self.score_nrow, self.score_ncol, self.score_row, self.score_col = model.model(self.input_images, is_training=True)
        self.variable_averages = tf.train.ExponentialMovingAverage(0.997, self.global_step)
        self.saver = tf.train.Saver(self.variable_averages.variables_to_restore())
        self.ckpt_state = tf.train.get_checkpoint_state(model_dir)
        print(self.ckpt_state)
        self.model_path = os.path.join(model_dir, os.path.basename(self.ckpt_state.model_checkpoint_path))
        self.saver.restore(self.session,self.model_path)


    def main_detection(self, image):
        im_resized, (ratio_h, ratio_w) = resize_image(image)
        # print('im_resize shape:', im_resized.shape)
        # cv2.imwrite('resize.jpg', im_resized)
        # im_resized = cv2.resize(im_resized, dsize=(512, 512), interpolation=cv2.INTER_AREA)
        # 将图像由BGR转化为RGB模式输入
        # im_resized = im_resized[:, :, ::-1].astype(np.float32)
        im_resized = im_resized[:, :, ::-1]
        score_nrow, score_ncol, score_row, score_col = self.session.run([self.score_nrow, self.score_ncol, self.score_row, self.score_col], feed_dict={self.input_images: [im_resized]})
        # plt.imshow(score_nrow[0,:,:,0])
        # plt.show()
        # breakpoint()
        # plt.imshow(score_ncol[0,:,:,0])
        # plt.imshow(score_row[0,:,:,0])
        # plt.imshow(score_col[0,:,:,0])
        # plt.show()
        return score_nrow[0], score_ncol[0], score_row[0], score_col[0], ratio_h, ratio_w

def resize_image(im):
    h, w, _ = im.shape
    size = (int(512), int(512))
    im = cv2.resize(im, size, interpolation=cv2.INTER_AREA)
    # la_p = cv2.resize(label_im, size, interpolation=cv2.INTER_AREA)

    ratio_h = 512 / float(h)
    ratio_w = 512 / float(w)

    return im, (ratio_h, ratio_w)

def dice_coefficient(y_true_cls, y_pred_cls):
    '''
    dice loss
    :param y_true_cls:
    :param y_pred_cls:
    :param training_mask:
    :return:
    '''
    eps = 1e-5
    intersection = np.sum(y_true_cls * y_pred_cls)
    union = np.sum(y_true_cls) + np.sum(y_pred_cls) + eps
    acc = 2 * intersection / union
    # tf.summary.scalar('classification_dice_loss', acc)
    return acc

if __name__ == '__main__':

    # 添加标签图像父路径
    gt_path = 'F:/laibo/Data/tijianbaogao/tip_bbg/test'

    result_path = './result/'
    instance = Detector('F:/laibo/Model/tableImageParser/tijianbaogao/ckpt/2021-11-17-512/')
    images = os.listdir(gt_path + '/image/')

    # 添加标签图像各子路径
    row_gt = gt_path + '/label_row'
    col_gt = gt_path + '/label_col'
    nrow_gt = gt_path + '/label_nrow'
    ncol_gt = gt_path + '/label_ncol'

    row_root = './tx_infer_data/row'
    col_root = './tx_infer_data/col'
    nrow_root = './tx_infer_data/nrow'
    ncol_root = './tx_infer_data/ncol'

    i_l = []
    total_acc = 0
    for x in range(len(images)):
        image_path = os.path.join(gt_path + '/image/', images[x])
        image_name = images[x]
        txt_name = image_name.replace('.jpg', '.txt')
        row_path = os.path.join(row_root, image_name)
        col_path = os.path.join(col_root, image_name)
        nrow_path = os.path.join(nrow_root, image_name)
        ncol_path = os.path.join(ncol_root, image_name)

        # 读取标签图像
        row_gt_image = cv2.imread('/'.join([row_gt, image_name]).replace('.jpg', '.png'), 0)
        col_gt_image = cv2.imread('/'.join([col_gt, image_name]).replace('.jpg', '.png'), 0)
        nrow_gt_image = cv2.imread('/'.join([nrow_gt, image_name]).replace('.jpg', '.png'), 0)
        ncol_gt_image = cv2.imread('/'.join([ncol_gt, image_name]).replace('.jpg', '.png'), 0)
        row_gt_image = row_gt_image[:, :, np.newaxis].astype(np.float32)
        col_gt_image = col_gt_image[:, :, np.newaxis].astype(np.float32)
        nrow_gt_image = nrow_gt_image[:, :, np.newaxis].astype(np.float32)
        ncol_gt_image = ncol_gt_image[:, :, np.newaxis].astype(np.float32)
        gt_nmap = cv2.bitwise_or(nrow_gt_image, ncol_gt_image)
        gt_lmap = cv2.bitwise_or(row_gt_image, col_gt_image)
        gt_map = cv2.bitwise_or(gt_nmap, gt_lmap)

        # image = cv2.imread(image_path, 0)
        image_color = cv2.resize(cv2.imread(image_path), (int(512), int(512)), interpolation=cv2.INTER_NEAREST)
        # instance.table_detection(image, image_color)
        score_nrow, score_ncol, score_row, score_col, ratio_h, ratio_w = instance.main_detection(image_color)

        score_nrow = np.where(score_nrow > 0.9, score_nrow, 0)
        score_nrow = np.where(score_nrow < 0.9, score_nrow, 1)

        score_ncol = np.where(score_ncol > 0.9, score_ncol, 0)
        score_ncol = np.where(score_ncol < 0.9, score_ncol, 1)

        score_row = np.where(score_row > 0.9, score_row, 0)
        score_row = np.where(score_row < 0.9, score_row, 1)

        score_col = np.where(score_col > 0.9, score_col, 0)
        score_col = np.where(score_col < 0.9, score_col, 1)

        nmap = cv2.bitwise_or(score_nrow, score_ncol)
        lmap = cv2.bitwise_or(score_row, score_col)
        pre_map = cv2.bitwise_or(nmap, lmap)

        # row_acc = dice_coefficient(row_gt_image, score_row)
        # col_acc = dice_coefficient(col_gt_image, score_col)
        # nrow_acc = dice_coefficient(nrow_gt_image, score_nrow)
        # ncol_acc = dice_coefficient(ncol_gt_image, score_ncol)
        acc = dice_coefficient(gt_map, pre_map)
        total_acc += acc
        print('image：', image_name)

        result = os.path.join(result_path, images[x])
        score_nrow_map = cv2.resize(score_nrow, dsize=None, fx=1/ratio_w, fy=1/ratio_h, interpolation=cv2.INTER_AREA)
        score_ncol_map = cv2.resize(score_ncol, dsize=None, fx=1 / ratio_w, fy=1 / ratio_h, interpolation=cv2.INTER_AREA)
        score_row_map = cv2.resize(score_row, dsize=None, fx=1 / ratio_w, fy=1 / ratio_h, interpolation=cv2.INTER_AREA)
        score_col_map = cv2.resize(score_col, dsize=None, fx=1 / ratio_w, fy=1 / ratio_h, interpolation=cv2.INTER_AREA)
        pre_map = cv2.resize(pre_map, dsize=None, fx=1 / ratio_w, fy=1 / ratio_h, interpolation=cv2.INTER_AREA)
        # mask_result = os.path.join(result_path, 'mask_'+images[x])
        # print(mask_result)

        cv2.imwrite(row_path, score_row*255)
        cv2.imwrite(col_path, score_col*255)
        cv2.imwrite(nrow_path, score_nrow * 255)
        cv2.imwrite(ncol_path, score_ncol * 255)
        cv2.imwrite(result, pre_map*255)

    print('acc:{}'.format(total_acc/len(images)))

